import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from itertools import count

# 定义马尔可夫链参数
states = ['A', 'B', 'C', 'D', 'E']
num_states = len(states)

# 每行是一个状态对其他状态的转移概率，必须保证每行之和为 1
transition_matrix = np.array([
    [0.2, 0.3, 0.2, 0.1, 0.2],
    [0.1, 0.4, 0.2, 0.2, 0.1],
    [0.3, 0.1, 0.1, 0.3, 0.2],
    [0.2, 0.2, 0.3, 0.1, 0.2],
    [0.1, 0.3, 0.1, 0.4, 0.1]
])

state_indices = {state: idx for idx, state in enumerate(states)}

# 初始化当前状态
current_state = np.random.choice(states)

# 构建有向图用于可视化
G = nx.DiGraph()

# 添加节点
for state in states:
    G.add_node(state)

# 添加边及其权重
edges = []
edge_labels = {}
for i, src in enumerate(states):
    for j, dst in enumerate(states):
        prob = transition_matrix[i][j]
        if prob > 0:
            edges.append((src, dst))
            edge_labels[(src, dst)] = f"{prob:.2f}"

G.add_edges_from(edges)

# 布局设置（使用圆形布局）
pos = nx.circular_layout(G)

# 路径记录
path = [current_state]

# 更新函数，用于动画每一帧
def update(frame):
    global current_state
    current_idx = state_indices[current_state]
    next_state = np.random.choice(states, p=transition_matrix[current_idx])
    path.append(next_state)
    current_state = next_state

    # 清空并重新绘制
    ax.clear()
    ax.set_title(f"Current State = {current_state}")

    # 绘制状态转移图
    nx.draw(G, pos, with_labels=True, node_size=800, node_color='lightblue', font_size=16, ax=ax, arrows=True)

    # 绘制边标签（转移概率）
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, ax=ax, font_size=10)

    # 高亮当前状态
    nx.draw_networkx_nodes(G, pos, nodelist=[current_state], node_color='orange', node_size=900, ax=ax)

    # 显示路径（显示最近10步）
    ax.text(0.5, -0.1, 'Path: ' + ' → '.join(path[-10:]), transform=ax.transAxes,
            ha='center', fontsize=10, color='black')

    return []

# 控制动画暂停/继续的函数
def on_key(event):
    if event.key == 'escape':
        ani.event_source.stop()
        print("Animation stopped by pressing ESC.")

# 创建绘图对象
fig, ax = plt.subplots(figsize=(8, 6))
ani = FuncAnimation(fig, update, frames=count(), interval=1000, blit=False)

# 绑定键盘事件
fig.canvas.mpl_connect('key_press_event', on_key)

plt.show()

# 输出最终路径
print("\nFinal Markov Chain Path:")
print(" -> ".join(path))
